from grid_env_ideal_obs_repeat_task_free_run_ctrl import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json

from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
from sklearn.pipeline import Pipeline
import matplotlib.pyplot as plt
import matplotlib as mpl

import matplotlib.animation as animation

from openTSNE import TSNE
from sklearn.manifold import Isomap

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# load landscape and states from file
def load_task(pth = "./logs/task.json"):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        print("state: ", state)
        print("goal: ", goal)
        print("landscape: ", landscape)
    return landscape, state, goal

event_type = ""
event_x = 0
event_y = 0

def rand_normal_like_tree(key: Any, params, std: float = 1.0, batch_shape: Tuple = ()):
    """Return a pytree like `params` where every element follows standard normal distribution
       May add a batch dim on parameters with batch_shape=(bs,)
    """
    num_vars = len(jax.tree_util.tree_leaves(params))
    treedef = jax.tree_util.tree_structure(params)

    all_keys = jax.random.split(key, num=num_vars)
    noise = jax.tree_util.tree_map(
        lambda g, k: std * jax.random.normal(k, shape= g.shape, dtype=g.dtype),
        params, jax.tree_util.tree_unflatten(treedef, all_keys))

    return noise

def main():

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration

    # cv2.namedWindow("rnn_state_img", 0)
    # cv2.namedWindow("img", cv2.WINDOW_GUI_NORMAL)

    k1 = jax.random.PRNGKey(npr.randint(0, 1000000))

    """ load model
    """
    params = load_weights(rpl_config.model_pth)

    key = jax.random.PRNGKey(npr.randint(0, 1000000))
    params_r = rand_normal_like_tree(key, params, std=1.0, batch_shape=(1,))

    """ create landscape
    """
    random_task = True
    # check if file on rpl_config.task_pth exists
    if os.path.isfile(rpl_config.task_pth):
        random_task = False

    if random_task:
        landscape = generate_maze_pool(num_mazes=1, width=10, height=10)
        landscape = padding_landscapes(landscape, width=12, height=12)
    else:
        landscape, state, goal = load_task(pth = rpl_config.task_pth)
        landscape = [landscape]

    print("landscape :")
    print(landscape)

    """ create agent
    """
    if rpl_config.nn_type == "vanilla":
        model = RNN(hidden_dims = rpl_config.nn_size)
    elif rpl_config.nn_type == "gru":
        model = GRU(hidden_dims = rpl_config.nn_size)

    # check if param fits the agent
    if rpl_config.nn_type == "vanilla":
        assert params["params"]["Dense_0"]["kernel"].shape[0] == rpl_config.nn_size + 10

    """ create grid env
    """
    start_time = time.time()
    GE = GridEnv(landscapes = landscape, width = 12, height = 12, num_envs_per_landscape = 1)
    GE.reset()
    print("time taken to create envs: ", time.time() - start_time)

    if not random_task:
        # set states of GE
        GE.batched_states = GE.batched_states.at[0, 0].set(state[0])
        GE.batched_states = GE.batched_states.at[0, 1].set(state[1])
        # set goals of GE
        GE.batched_goals = GE.batched_goals.at[0, 0].set(goal[0])
        GE.batched_goals = GE.batched_goals.at[0, 1].set(goal[1])
        GE.init_batched_states, GE.init_batched_goals = jnp.copy(GE.batched_states), jnp.copy(GE.batched_goals)
        GE.batched_goal_reached = batch_compute_goal_reached(GE.batched_states, GE.batched_goals)
        GE.last_batched_goal_reached = jnp.copy(GE.batched_goal_reached)
        GE.concat_obs = get_ideal_obs_vmap(GE.batched_envs, GE.batched_states, GE.batched_goals, GE.last_batched_goal_reached)

    concat_obs = GE.concat_obs

    rnn_state = model.initial_state(GE.num_envs)
    step_count = 0

    rnn_state_waterfall = []
    rnn_state_waterfall_new = []

    goal_signal = []

    """ rnn state visualization
    """
    neural_populations = [1, 5, 6, 11, 13, 14, 15, 16, 17, 18, 19, 22, 23, 24, 25, 27, 28, 34, 44, 46, 57, 62, 65, 68, 80, 82, 83, 88, 89, 96, 98, 100, 102, 104, 108, 113, 116, 120, 122]
    pop_centers =[[0.9999582767486572, -0.9831897020339966, 0.9996467232704163, -0.9994123578071594, -0.999851405620575, -0.9999927878379822, 0.999983549118042, 0.9612636566162109, 0.9998806715011597, 0.9932894706726074, -0.9997431635856628, 0.9999991059303284, 0.9999998807907104, -0.9999634027481079, 0.9999784231185913, 0.9996755123138428, -0.9992828369140625, 0.8952787518501282, -0.9999989867210388, 0.9748570919036865, 0.9935382008552551, -0.9999995827674866, 0.9871329665184021, 0.9740328788757324, -0.9953449368476868, 0.9701042771339417, -0.9991737604141235, -0.9999988675117493, -0.9990499019622803, -0.999997079372406, 0.9949736595153809, 0.9822999238967896, -0.9999629855155945, -0.9999940991401672, -0.9986229538917542, 0.9996981024742126, 0.9999974370002747, 0.9871447682380676, -0.9999998807907104]
    ,[-0.9998412728309631, 0.9671010375022888, -0.9999982714653015, 0.9992090463638306, 0.9994693994522095, 0.9999602437019348, -0.9925903081893921, -0.9984927773475647, -0.9941282868385315, 0.9965940713882446, 0.9773151278495789, -0.9977096915245056, -0.9996675252914429, 0.9999555349349976, -0.9999988675117493, -0.9978929758071899, 0.9999995231628418, -0.9997955560684204, 0.9996235966682434, 0.9973439574241638, -0.9999579787254333, 0.9999728798866272, 0.9823980927467346, -0.9997960329055786, 0.9941484332084656, -0.9999427795410156, 0.9995871186256409, 0.9999239444732666, 0.9942917823791504, 0.9999849796295166, -0.9987909197807312, -0.9999961256980896, 0.9863384366035461, 0.9999999403953552, 0.9593451619148254, -0.9977843165397644, -0.9793924689292908, -0.9543303847312927, 0.9943166375160217]
    ,[0.9999920725822449, -0.9973777532577515, 0.9999567866325378, -0.9910239577293396, -0.9999980926513672, -0.9999908804893494, 0.9994378089904785, 0.9657203555107117, 0.9999774098396301, -0.9924333691596985, -0.9999697804450989, 0.9999998807907104, 0.9999997615814209, -0.9969027638435364, 0.9999698996543884, 0.9999940991401672, -0.999999463558197, 0.9991713166236877, -0.9999998807907104, -0.9957963228225708, 0.9988913536071777, -0.9999998807907104, -0.9945865869522095, 0.9999521970748901, -0.9996950626373291, 0.9985164403915405, -0.998382568359375, -0.9999978542327881, -0.9999244809150696, -0.9996201395988464, 0.999308168888092, 0.9999228715896606, -0.9999930262565613, -0.9999990463256836, -0.9861818552017212, 0.9988118410110474, 0.9999998807907104, 0.9948789477348328, -0.9999997615814209]
    ,[-0.9705339074134827, 0.9958152174949646, -0.9995390772819519, 0.9999624490737915, 0.999870240688324, 0.9999515414237976, -0.9939287304878235, -0.9997881054878235, -0.9936801791191101, -0.977760910987854, 0.9972586035728455, -0.9999999403953552, -0.9999999403953552, 0.9999993443489075, -0.9999997019767761, -0.9981520771980286, 0.994085431098938, -0.9939032793045044, 0.9999999403953552, -0.991743266582489, -0.9992260932922363, 0.9999999403953552, -0.9556638598442078, -0.9771517515182495, 0.9983199238777161, -0.9999507069587708, 0.9929742217063904, 0.9999997615814209, 0.9999983906745911, 0.9997593760490417, -0.9996225833892822, -0.9998718500137329, 0.9999980330467224, 0.9999902844429016, 0.9994619488716125, -0.9978642463684082, -0.9999993443489075, -0.9995915293693542, 0.9999999403953552]]


    neural_populations = [0, 1, 3, 4, 5, 6, 7, 8, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 22, 23, 24, 25, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 39, 41, 43, 44, 46, 49, 52, 54, 57, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 71, 74, 77, 79, 80, 82, 83, 84, 85, 86, 88, 89, 90, 91, 92, 96, 97, 98, 99, 100, 101, 102, 103, 104, 106, 107, 108, 109, 110, 113, 115, 116, 118, 120, 121, 122, 123, 124, 125, 126]
    pop_centers =[[-0.9461527466773987, 0.9999582767486572, -0.975846529006958, -0.9906558990478516, -0.9831897020339966, 0.9996467232704163, 0.9916031956672668, 0.9968320727348328, -0.9124768972396851, -0.9994123578071594, -0.8605355620384216, -0.999851405620575, -0.9999927878379822, 0.999983549118042, 0.9612636566162109, 0.9998806715011597, 0.9932894706726074, -0.9997431635856628, -0.9999997019767761, 0.9999991059303284, 0.9999998807907104, -0.9999634027481079, 0.9999784231185913, 0.9996755123138428, -0.9992828369140625, 0.8932662010192871, -0.9987356662750244, 0.9904150366783142, -0.999648928642273, 0.9468201994895935, 0.8952787518501282, -0.9989546537399292, -0.9982856512069702, 0.9999709129333496, -0.9981957077980042, -0.9974076151847839, -0.9678183197975159, -0.9999989867210388, 0.9748570919036865, -0.9999744892120361, 0.9587243795394897, 0.9830151200294495, 0.9935382008552551, -0.9996662139892578, -0.999999463558197, -0.9789413213729858, -0.9999995827674866, -0.8860974311828613, 0.9345378875732422, 0.9871329665184021, 0.8611211180686951, -0.9988968968391418, 0.9740328788757324, 0.995445966720581, 0.9981663227081299, -0.9995632767677307, 0.9999988675117493, -0.8693735599517822, -0.9953449368476868, 0.9701042771339417, -0.9991737604141235, -0.9907732009887695, -0.990005612373352, 0.9999998807907104, -0.9999988675117493, -0.9990499019622803, -0.946376621723175, 0.9362999200820923, -0.9999972581863403, -0.999997079372406, 0.9736419916152954, 0.9949736595153809, -0.964794933795929, 0.9822999238967896, -0.970579206943512, -0.9999629855155945, -0.999987781047821, -0.9999940991401672, -0.9361556768417358, -0.9990412592887878, -0.9986229538917542, -0.9984415769577026, 0.9999998211860657, 0.9996981024742126, -0.9994224309921265, 0.9999974370002747, 0.9792995452880859, 0.9871447682380676, 0.9999986886978149, -0.9999998807907104, 0.9999998807907104, 0.9490108489990234, 0.9829628467559814, -0.9778906106948853]
    ,[-0.9283203482627869, -0.9998435378074646, -0.9638441205024719, 0.9988398551940918, 0.967153012752533, -0.9999982714653015, -0.9427096843719482, -0.9438421130180359, 0.9670294523239136, 0.999212384223938, -0.9875056743621826, 0.9994784593582153, 0.9999609589576721, -0.9926629662513733, -0.998508095741272, -0.9942294359207153, 0.9966515898704529, 0.9777061939239502, 0.9663264751434326, -0.9977492094039917, -0.9996732473373413, 0.9999483227729797, -0.9999988675117493, -0.9979292154312134, 0.9999995231628418, -0.9752395749092102, 0.922939658164978, 0.9681260585784912, -0.9128158092498779, -0.9994750022888184, -0.9997990131378174, 0.917654275894165, 0.9677172303199768, -0.9653894305229187, -0.8428236842155457, 0.9987689256668091, -0.9331029057502747, 0.9996300935745239, 0.9973886609077454, 0.9643940925598145, 0.9969982504844666, -0.9655174612998962, -0.9999585151672363, 0.965191125869751, 0.9669886231422424, 0.9599414467811584, 0.9999733567237854, 0.9693759679794312, 0.9535602331161499, 0.9827015399932861, 0.9561278820037842, 0.9657979011535645, -0.9997995495796204, -0.9655174016952515, -0.9558514356613159, 0.9520189762115479, -0.9576465487480164, -0.9739927053451538, 0.9942493438720703, -0.9999437928199768, 0.9995939135551453, 0.9995091557502747, 0.9656971096992493, -0.9655343890190125, 0.9999252557754517, 0.9943885207176208, 0.9999865889549255, -0.9767569899559021, 0.9655064940452576, 0.9999851584434509, 0.9700678586959839, -0.9988107085227966, -0.9514037370681763, -0.9999961853027344, 0.958463728427887, 0.9865739941596985, 0.8932945728302002, 0.9999999403953552, 0.9678045511245728, 0.9656780362129211, 0.9585627317428589, 0.9654185771942139, -0.9651579260826111, -0.9978225231170654, 0.965463399887085, -0.9797477722167969, -0.9656004309654236, -0.9550648331642151, -0.9615163207054138, 0.9944142699241638, -0.9655593633651733, -0.9863012433052063, 0.8636111617088318, 0.9887285828590393]
    ,[0.9958351254463196, 0.9999920725822449, 0.9943109154701233, -0.9985464215278625, -0.9973777532577515, 0.9999567866325378, 0.9958251118659973, 0.9876009225845337, -0.9994816184043884, -0.9910239577293396, 0.9949482679367065, -0.9999980926513672, -0.9999908804893494, 0.9994378089904785, 0.9657203555107117, 0.9999774098396301, -0.9924333691596985, -0.9999697804450989, -0.9962242245674133, 0.9999998807907104, 0.9999997615814209, -0.9969027638435364, 0.9999698996543884, 0.9999940991401672, -0.999999463558197, 0.9914127588272095, -0.9996015429496765, -0.9995436072349548, 0.97539883852005, 0.9999948143959045, 0.9991713166236877, -0.9886994957923889, -0.9960036873817444, 0.9949798583984375, 0.9921016693115234, -0.9937861561775208, 0.9988359212875366, -0.9999998807907104, -0.9957963228225708, -0.9823622107505798, -0.9970131516456604, 0.9999969005584717, 0.9988913536071777, -0.9997882843017578, -0.9999831318855286, -0.9045405983924866, -0.9999998807907104, -0.9996019005775452, -0.9999549984931946, -0.9945865869522095, -0.9996899962425232, -0.9977795481681824, 0.9999521970748901, 0.9999998807907104, 0.9997276067733765, -0.692555844783783, 0.995883047580719, 0.9771542549133301, -0.9996950626373291, 0.9985164403915405, -0.998382568359375, -0.9999997615814209, -0.9999960064888, 0.9999998807907104, -0.9999978542327881, -0.9999244809150696, -0.9996898770332336, 0.9799578189849854, -0.999114990234375, -0.9996201395988464, -0.9984035491943359, 0.999308168888092, 0.9999555349349976, 0.9999228715896606, -0.9461102485656738, -0.9999930262565613, -0.9784282445907593, -0.9999990463256836, -0.9999871253967285, -0.9881334900856018, -0.9861818552017212, -0.9999997019767761, 0.9999986886978149, 0.9988118410110474, -0.9999892711639404, 0.9999998807907104, 0.9997080564498901, 0.9948789477348328, 0.9999954700469971, -0.9999997615814209, 0.9999998807907104, 0.9999998807907104, -0.917002260684967, -0.9996073842048645]
    ,[0.929304838180542, -0.9705339074134827, 0.9548043012619019, 0.9257931709289551, 0.9958152174949646, -0.9995390772819519, -0.9857621788978577, -0.9858974814414978, 0.9543991088867188, 0.9999624490737915, 0.990217387676239, 0.999870240688324, 0.9999515414237976, -0.9939287304878235, -0.9997881054878235, -0.9936801791191101, -0.977760910987854, 0.9972586035728455, 1.0, -0.9999999403953552, -0.9999999403953552, 0.9999993443489075, -0.9999997019767761, -0.9981520771980286, 0.994085431098938, -0.9613664746284485, 0.9221969842910767, -0.9750844836235046, 0.9994914531707764, -0.9958109855651855, -0.9939032793045044, 0.9942757487297058, 0.9563018083572388, -0.9999991059303284, 0.995768666267395, 0.9596813917160034, 0.999951958656311, 0.9999999403953552, -0.991743266582489, 0.9800897836685181, -0.8371233940124512, -0.9999257922172546, -0.9992260932922363, 0.9999868273735046, 0.9999557137489319, 0.9107763171195984, 0.9999999403953552, 0.997278094291687, -0.9970739483833313, -0.9556638598442078, -0.9633026123046875, 0.9999330639839172, -0.9771517515182495, -0.9751230478286743, -0.9946858286857605, 0.9999380707740784, -0.9999024271965027, 0.8732730746269226, 0.9983199238777161, -0.9999507069587708, 0.9929742217063904, 0.8123788833618164, 0.9945662617683411, -0.9999999403953552, 0.9999997615814209, 0.9999983906745911, 0.9869656562805176, -0.9268302917480469, 0.9999990463256836, 0.9997593760490417, -0.98912513256073, -0.9996225833892822, 0.9998481273651123, -0.9998718500137329, 0.9972803592681885, 0.9999980330467224, 0.9946923851966858, 0.9999902844429016, 0.9754811525344849, 0.9999573826789856, 0.9994619488716125, 0.9924792051315308, -0.9999999403953552, -0.9978642463684082, 0.9987314343452454, -0.9999993443489075, -0.9999983906745911, -0.9995915293693542, -0.9999723434448242, 0.9999999403953552, -0.9999970197677612, -0.9448661804199219, -0.8927628397941589, 0.90985506772995]]

    pop_centers =  [[ 0.99583536,  0.9999922 ,  0.994311  , -0.9985466 , -0.99737775,
    0.99995697,  0.9958252 ,  0.9876008 , -0.99948156, -0.9910242 ,
    0.99494815, -0.9999976 , -0.9999907 ,  0.9994379 ,  0.9657203 ,
    0.99997747, -0.99243313, -0.9999696 , -0.9962243 ,  0.9999992 ,
    1.0000002 , -0.99690264,  0.99996954,  0.99999356, -0.99999976,
    0.9914125 , -0.9996015 , -0.99954337,  0.97539896,  0.9999945 ,
    0.99917144, -0.98869944, -0.99600375,  0.99498   ,  0.99210197,
    -0.9937862 ,  0.9988358 , -0.99999946, -0.99579597, -0.9823625 ,
    -0.9970135 ,  0.999997  ,  0.9988911 , -0.99978805, -0.99998295,
    -0.9045407 , -0.99999976, -0.99960166, -0.9999547 , -0.99458677,
    -0.99968934, -0.99777937,  0.9999517 ,  0.99999917,  0.9997276 ,
    -0.6925559 ,  0.9958831 ,  0.97715414, -0.9996948 ,  0.9985169 ,
    -0.99838257, -0.99999994, -0.9999956 ,  1.0000002 , -0.9999974 ,
    -0.99992454, -0.9996894 ,  0.9799579 , -0.999115  , -0.9996197 ,
    -0.9984035 ,  0.9993084 ,  0.99995553,  0.99992245, -0.9461104 ,
    -0.9999931 , -0.97842836, -0.9999994 , -0.9999875 , -0.9881337 ,
    -0.9861822 , -1.        ,  0.9999982 ,  0.99881184, -0.99998903,
    0.9999999 ,  0.9997082 ,  0.9948789 ,  0.9999956 , -0.99999934,
    0.99999917,  1.0000007 , -0.91700256, -0.9996074 ],
    [-0.9621505 , -0.9998406 , -0.9982953 ,  0.9988196 ,  0.96657693,
    -0.9999994 , -0.94170475, -0.9778629 ,  0.99976325,  0.99935323,
    -0.9872892 ,  0.9997662 ,  0.9999614 , -0.9989443 , -0.99848205,
    -0.9999448 ,  0.99659777,  0.99903643,  0.99981993, -1.0000005 ,
    -0.99999964,  0.9999475 , -0.9999983 , -0.99925727,  0.99999887,
    -0.985652  ,  0.95659566,  0.9998184 , -0.9460587 , -0.99999726,
    -0.99979585,  0.91629106,  0.9984362 , -0.99957293, -0.8400663 ,
    0.9987501 , -0.956525  ,  1.0000004 ,  0.9994299 ,  0.9988571 ,
    0.99972856, -0.99999964, -0.9999803 ,  0.9996511 ,  0.99906147,
    0.9592386 ,  1.0000002 ,  0.99979794,  0.9878042 ,  0.9845369 ,
    0.9553718 ,  0.99936366, -0.9998006 , -1.0000005 , -0.98865604,
    0.98555565, -0.9917041 , -0.9735364 ,  0.99414843, -0.99994516,
    0.9996034 ,  0.999998  ,  0.99996114, -0.9999998 ,  0.9999818 ,
    0.999977  ,  0.9999913 , -0.9978424 ,  0.9999783 ,  0.9999853 ,
    0.9970532 , -0.9988095 , -0.9855086 , -0.99999636,  0.99270505,
    0.99999356,  0.8930451 ,  0.99999917,  0.99997294,  0.99984807,
    0.9578357 ,  0.9998978 , -0.9996089 , -0.9979516 ,  0.99994546,
    -0.99999994, -0.9999998 , -0.95427656, -0.98555833,  0.9998825 ,
    -1.0000002 , -0.99999917,  0.89630413,  0.98853105],
    [-0.94615287,  0.99995846, -0.9758465 , -0.99065584, -0.9831897 ,
    0.9996464 ,  0.99160326,  0.9968322 , -0.9124768 , -0.9994121 ,
    -0.8605355 , -0.99985135, -0.9999925 ,  0.9999834 ,  0.9612637 ,
    0.99988055,  0.9932893 , -0.9997432 , -1.        ,  0.9999985 ,
    1.0000004 , -0.9999632 ,  0.9999779 ,  0.9996752 , -0.9992828 ,
    0.89326644, -0.9987357 ,  0.99041516, -0.9996488 ,  0.94682026,
    0.8952788 , -0.99895436, -0.9982858 ,  0.99997085, -0.9981955 ,
    -0.9974078 , -0.9678181 , -0.99999857,  0.97485703, -0.9999744 ,
    0.9587243 ,  0.9830153 ,  0.9935382 , -0.99966586, -0.99999917,
    -0.97894144, -0.99999946, -0.8860976 ,  0.9345378 ,  0.98713297,
    0.86112106, -0.9988967 ,  0.9740331 ,  0.99544585,  0.9981664 ,
    -0.9995633 ,  0.99999917, -0.8693738 , -0.99534494,  0.9701041 ,
    -0.99917364, -0.99077314, -0.9900056 ,  1.0000002 , -0.99999833,
    -0.9990499 , -0.94637686,  0.93629986, -0.99999714, -0.9999965 ,
    0.97364193,  0.9949735 , -0.9647948 ,  0.9823    , -0.9705795 ,
    -0.99996287, -0.9999879 , -0.99999446, -0.93615556, -0.9990414 ,
    -0.99862295, -0.9984418 ,  0.99999934,  0.99969774, -0.9994222 ,
    0.9999975 ,  0.97929955,  0.9871444 ,  0.9999993 , -0.9999995 ,
    0.9999993 ,  0.94901055,  0.982963  , -0.9778905 ],
    [ 0.93077755, -0.9711479 ,  0.95574296,  0.9273392 ,  0.9959022 ,
    -0.9995469 , -0.9860587 , -0.9446222 ,  0.9157908 ,  0.99977916,
    0.9487575 ,  0.9995199 ,  0.99995184, -0.98644364, -0.99979264,
    -0.9869046 , -0.9365635 ,  0.971522  ,  0.95952445, -0.99728096,
    -0.99960494,  0.9999996 , -0.9999991 , -0.99657017,  0.9942079 ,
    -0.94929045,  0.88224614, -0.97223574,  0.99912745, -0.99526775,
    -0.9940303 ,  0.9942979 ,  0.920061  , -0.9586844 ,  0.9541901 ,
    0.9605179 ,  0.98749375,  0.99955344, -0.95272684,  0.93883806,
    -0.80215514, -0.9582608 , -0.99921507,  0.9583404 ,  0.9611827 ,
    0.91263527,  0.9999681 ,  0.9605707 , -0.99710023, -0.91746074,
    -0.9224168 ,  0.9593625 , -0.97762173, -0.9339752 , -0.95492136,
    0.959115  , -0.95857835,  0.83424675,  0.99835485, -0.99994755,
    0.99310076,  0.8156971 ,  0.953277  , -0.95835394,  0.9999319 ,
    0.9932449 ,  0.98723185, -0.9028315 ,  0.95834553,  0.9997646 ,
    -0.980354  , -0.99960726,  0.99969685, -0.99987465,  0.95581   ,
    0.9837817 ,  0.99287623,  0.99998987,  0.9371209 ,  0.95866543,
    0.9994732 ,  0.95097154, -0.9583645 , -0.9977102 ,  0.9570914 ,
    -0.9755282 , -0.9584324 , -0.9995997 , -0.97062093,  0.9933904 ,
    -0.95838153, -0.9294621 , -0.8949945 ,  0.9117331 ]]
    neural_populations =  [0, 1, 3, 4, 5, 6, 7, 8, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 22, 23, 24, 25, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 39, 41, 43, 44, 46, 49, 52, 54, 57, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 71, 74, 77, 79, 80, 82, 83, 84, 85, 86, 88, 89, 90, 91, 92, 96, 97, 98, 99, 100, 101, 102, 103, 104, 106, 107, 108, 109, 110, 113, 115, 116, 118, 120, 121, 122, 123, 124, 125, 126]

    pca = PCA()
    pipe = Pipeline([('scaler', StandardScaler()), ('pca', pca)])

    quadrant_beleived_old = 0
    quadrant_list = []

    quadrant_beleived_old1 = 0
    quadrant_list1 = []

    test_period = rpl_config.life_duration
    goal_appear_time = test_period//2
    # goal_appear_time = test_period**2
    goal_arrived = False

    # get original trajectory
    for t_ in range(test_period * 2):

        step_count += 1

        """ model forward and step the env
        """
        rnn_state, y1 = model_forward(params, rnn_state, concat_obs, model)
        # rnn_state, y1 = model_forward(params_r, rnn_state, concat_obs, model)
        batched_actions = get_action_vmap(y1)

        if t_ < goal_appear_time:
            batched_goal_reached, concat_obs = GE.step(batched_actions, reset = False)
            goal_signal.append(0)
        elif t_ >= goal_appear_time and t_ <= test_period:
            # # test1
            # batched_goal_reached, concat_obs = GE.step(batched_actions, reset = True)
            # if batched_goal_reached[0] == 1:
            #     goal_signal.append(1)
            # else:
            #     goal_signal.append(0)

            # # test2
            # batched_goal_reached, concat_obs = GE.step(batched_actions, reset = False)
            # goal_signal.append(1)
            # concat_obs = concat_obs.at[0,4].set(-2)

            # test3
            batched_goal_reached, concat_obs = GE.step(batched_actions, reset = False)
            # set the elements of rnn_state[0][var_clt_mean_less_than_0_1] to rnn_feature_3
            # 将 rnn_state[0] 转换为 numpy 数组
            rnn_state_0 = np.array(rnn_state[0].copy())
            # 遍历 var_clt_mean_less_than_0_1 中的每个元素，将其替换为 rnn_feature_3 中对应的值
            for i in range(len(neural_populations)):
                rnn_state_0[neural_populations[i]] = pop_centers[3][i]
            rnn_state = rnn_state.at[0].set(jnp.array(rnn_state_0))
            goal_signal.append(1)
            rnn_state_waterfall_new.append(rnn_state[0].tolist())

        else:
            batched_goal_reached, concat_obs = GE.step(batched_actions, reset = False)
            goal_signal.append(0)
        
        rnn_state_waterfall.append(rnn_state[0].tolist())
        

        # compute the variance of rnn_state_waterfall in a window of 20 steps
        rnn_state_window = rnn_state_waterfall[-20:]
        rnn_state_window_var = np.var(rnn_state_window, axis=0)

        # print("rnn_state_window_var: ", rnn_state_window_var)
        # print("")

        static_feature = rnn_state[0, neural_populations]
        sim_to_feature0 = np.dot(static_feature, pop_centers[0])
        sim_to_feature1 = np.dot(static_feature, pop_centers[1])
        sim_to_feature2 = np.dot(static_feature, pop_centers[2])
        sim_to_feature3 = np.dot(static_feature, pop_centers[3])
        quadrant_beleived = np.argmax([sim_to_feature0, sim_to_feature1, sim_to_feature2, sim_to_feature3])

        quadrant_list.append(quadrant_beleived)

    print("quadrant_list :", quadrant_list)
    print("goal_signal len:", len(goal_signal))

    # plot quadrant_list
    plt.figure()
    plt.plot(quadrant_list)
    plt.plot(goal_signal)
    plt.show()

    # do PCA on rnn_state_waterfall
    def colorFader(c1,c2,mix=0): #fade (linear interpolate) from color c1 (at mix=0) to c2 (mix=1)
        c1=np.array(mpl.colors.to_rgb(c1))
        c2=np.array(mpl.colors.to_rgb(c2))
        return mpl.colors.to_hex((1-mix)*c1 + mix*c2)
    c1='red' #blue
    c2='blue' #green

    pca = PCA()
    pca.fit(np.array(rnn_state_waterfall))
    pca_components = pca.components_[:5]
    print("shape of pca_components: ", pca_components.shape)
    # print(np.array2string(pca_components, separator=','))

    xt = pca.transform(np.array(rnn_state_waterfall))

    k1 = 0
    k2 = 1
    k3 = 2

    fig = plt.figure()
    ax = plt.axes(projection='3d')
    n = len(rnn_state_waterfall)-1
    for i in range(0,n):
        ax.plot([xt[i, k1], xt[i+1, k1]], [xt[i, k2], xt[i+1, k2]], [xt[i, k3], xt[i+1, k3]], color=colorFader(c1,c2,i/n), linewidth = 1)
    line, = ax.plot([], [], [], 'ro', markersize=20)
    # set line color
    line.set_color('red')
    line.set_markevery([])

    # 动画更新函数 
    def animate(i): 
        line.set_data(xt[i:i+1, k1], xt[i:i+1, k2])
        line.set_3d_properties(xt[i:i+1, k3])
        line.set_markevery([0])
        if goal_signal[i] == 1:
            line.set_color('green')
            # print("goal_signal[i] == 1 ",i)
        else:
            line.set_color('red')
        return line,

    print("xt.shape[0]: ", xt.shape[0])
    ani = animation.FuncAnimation(fig, animate, frames=xt.shape[0], interval=10, blit=False)

    plt.show()


    # xt2
    pca.fit(np.array(rnn_state_waterfall_new))
    xt2 = pca.transform(np.array(rnn_state_waterfall_new))

    pca_components = pca.components_[:5]
    print("shape of pca_components: ", pca_components.shape)
    print(np.array2string(pca_components, separator=','))

    # plot the first three dimensions of xt2
    fig = plt.figure()
    ax = plt.axes(projection='3d')
    n = len(rnn_state_waterfall_new)-1
    for i in range(0,n):
        ax.plot([xt2[i, k1], xt2[i+1, k1]], [xt2[i, k2], xt2[i+1, k2]], [xt2[i, k3], xt2[i+1, k3]], color=colorFader(c1,c2,i/n), linewidth = 1)
    line, = ax.plot([], [], [], 'ro', markersize=20)
    # set line color
    line.set_color('red')
    line.set_markevery([])

    # 动画更新函数 
    def animate(i): 
        line.set_data(xt2[i:i+1, k1], xt2[i:i+1, k2])
        line.set_3d_properties(xt2[i:i+1, k3])
        line.set_markevery([0])
        return line,
    ani = animation.FuncAnimation(fig, animate, frames=xt2.shape[0], interval=1, blit=False)

    plt.show()




if __name__ == "__main__":
    main()